#pragma once

#include "my_defs.h"
#include "my_atomic.h"
#include "my_spinlock.h"

#define MY_MEM_ITEM_MAGIC                               0xFA5CFFFA
#define MY_MEM_ITEM_RAW_MAGIC                           0xFA5CFFFB
#define MY_ALIGNED_MEM_ITEM_MAGIC                       0xFA5CFFAF
#define MY_ALIGNED_MEM_ITEM_RAW_MAGIC                   0xFA5CFFBF
#define MY_MEM_4K_ALIGNMENT                             (4 * 1024)

typedef void *(*my_malloc_pt)(int64_t size);
typedef void (*my_free_pt)(void *ptr);

typedef void *(*my_aligned_malloc_pt)(int64_t alignment, int64_t size);
typedef void (*my_aligned_free_pt)(void *ptr);

#ifdef MEM_DEBUG
extern my_atomic_int64_t g_mem_hold_num;
#endif

static inline void my_mem_assert_leak()
{
#ifdef MEM_DEBUG
    assert(my_atomic_int64_load(&g_mem_hold_num) == 0);
#endif
}

static inline void *
my_malloc(int64_t size)
{
    void *ptr = malloc((size_t)size);
    assert(ptr != NULL);

#ifdef MEM_DEBUG
    my_atomic_int64_add_fetch(&g_mem_hold_num, 1);
#endif

    return ptr;
}

static inline void *
my_calloc(int64_t size)
{
    void *ptr = my_malloc(size);
    bzero(ptr, (size_t)size);
    return ptr;
}

static inline void
my_free(void *ptr)
{
    free(ptr);
#ifdef MEM_DEBUG
    my_atomic_int64_add_fetch(&g_mem_hold_num, -1);
#endif 
}

typedef struct {
    uint32_t            magic;
    uint32_t            size;
    void                *next;  
} my_mem_item_head_t;

static inline void *
my_mem_raw_alloc(int64_t size)
{
    my_mem_item_head_t *item = NULL;
    item = my_malloc((int64_t)sizeof(my_mem_item_head_t) + size);
    item->magic = MY_MEM_ITEM_RAW_MAGIC;
    /* if size > 4G, the size will be truncated */
    item->size = (uint32_t)size;
    item->next = NULL;
    return (void *)item + sizeof(my_mem_item_head_t);
}

static inline void *
my_mem_alloc(int64_t size)
{
    my_mem_item_head_t *item = NULL;
    item = my_malloc((int64_t)sizeof(my_mem_item_head_t) + size);
    item->magic = MY_MEM_ITEM_MAGIC;
    /* if size > 4G, the size will be truncated */
    item->size = (uint32_t)size;
    item->next = NULL;
    return (void *)item + sizeof(my_mem_item_head_t);
}

static inline void
my_mem_free(void *ptr)
{
    my_mem_item_head_t *item;
    item = ptr - sizeof(my_mem_item_head_t);

    if (item->magic != MY_MEM_ITEM_RAW_MAGIC &&
        item->magic != MY_MEM_ITEM_MAGIC) {
        assert(!"unexcept item magic");
    }

    my_free(item);
}

typedef struct {
    uint32_t            magic;
    uint32_t            size;
    void                *next; /* for free list */
    void                *ptr;  /* p2->p1 */
} my_aligned_mem_item_header_t;

static inline void *
my_aligned_raw_malloc(int64_t alignment, int64_t size)
{
    int64_t offset = (alignment - 1) + (int64_t)sizeof(my_aligned_mem_item_header_t);
    my_aligned_mem_item_header_t *item = my_malloc(size + offset);
    item->magic = MY_ALIGNED_MEM_ITEM_RAW_MAGIC;
    item->next = item->ptr = NULL;
    item->size = (uint32_t)size;
    void *p1 = (void *)item;
    void **p2 = (void**)(((size_t)p1 + (size_t)offset) & ~((size_t)alignment - 1));
    p2[-1] = p1;
    return p2;
}

static inline void *
my_aligned_malloc(int64_t alignment, int64_t size)
{
    int64_t offset = (alignment - 1) + (int64_t)sizeof(my_aligned_mem_item_header_t);
    my_aligned_mem_item_header_t *item = my_malloc(size + offset);
    item->magic = MY_ALIGNED_MEM_ITEM_MAGIC;
    item->next = item->ptr = NULL;
    item->size = (uint32_t)size;
    void *p1 = (void *)item;
    void **p2 = (void**)(((size_t)p1 + (size_t)offset) & ~((size_t)alignment - 1));
    p2[-1] = p1;
    return p2;
}

static inline void
my_aligned_free(void *ptr)
{
    void *p1 = ((void **)ptr)[-1];
    my_aligned_mem_item_header_t *item = p1;


    if (item->magic != MY_ALIGNED_MEM_ITEM_RAW_MAGIC &&
        item->magic != MY_ALIGNED_MEM_ITEM_MAGIC) {
        assert(!"unexcept item magic");
    }

    my_free(item);
}


typedef struct {
    my_spinlock_t           lock;
    uint32_t                space_size;
    int64_t                 water_level;
    int64_t                 count;
    my_mem_item_head_t      head;
    my_malloc_pt            malloc_pt;
    my_free_pt              free_pt;
} my_mem_freelist_t;

static inline void
my_mem_freelist_init(my_mem_freelist_t *list, uint32_t space_size,
    int64_t water_level, my_malloc_pt malloc_pt, my_free_pt free_pt)
{

    if (!malloc_pt || !free_pt) {
        assert(!"malloc or free handler is NULL");
    }

    my_spinlock_init(&list->lock);
    list->space_size = space_size;
    list->water_level = water_level;
    list->count = 0;
    list->head.magic = MY_MEM_ITEM_MAGIC;
    list->head.size = space_size;
    list->head.next = NULL;
    list->malloc_pt = malloc_pt;
    list->free_pt = free_pt;
}

static inline void *
my_mem_freelist_alloc(my_mem_freelist_t *list)
{
    my_mem_item_head_t *item = NULL;

    my_spinlock_lock(&list->lock);
    if (list->count > 0) {
        item = list->head.next;
        list->head.next = item->next;
        --list->count;
    }
    my_spinlock_unlock(&list->lock);

    if (item != NULL) {
        item->next = NULL;
        return (void *)item + sizeof(my_mem_item_head_t);
    }

    void *ptr = list->malloc_pt(list->space_size);
    // item = ptr = sizeof(my_mem_item_head_t);
    // if (item->magic != MY_MEM_ITEM_MAGIC &&
    //     item->magic != MY_MEM_ITEM_RAW_MAGIC) {
    //     assert(!"unexcept item magic");
    // }

    return ptr;
}

static inline void *
my_mem_freelist_calloc(my_mem_freelist_t *list)
{
    void *ptr = my_mem_freelist_alloc(list);
    bzero(ptr, (size_t)list->space_size);
    return ptr;
}

static inline void
my_mem_freelist_free(my_mem_freelist_t *list, void *ptr)
{
    my_mem_item_head_t *item;
    item = ptr - sizeof(my_mem_item_head_t);

    if (item->magic != MY_MEM_ITEM_MAGIC) {
        assert(!"unexcept magic number");
    }

    if (item->size != list->space_size) {
        assert(!"unexcept stats");
    }

    my_spinlock_lock(&list->lock);
    if (list->count < list->water_level) {
        item->next = list->head.next;
        list->head.next = item;
        item = NULL;
        ++list->count;
    }
    my_spinlock_unlock(&list->lock);

    if (item == NULL) {
        return;
    }

    list->free_pt(ptr);
}

static inline void
my_mem_freelist_clean(my_mem_freelist_t *list)
{
    void                *ptr;
    my_mem_item_head_t  *item;

    my_spinlock_lock(&list->lock);
    while(list->head.next != NULL) {
        item = list->head.next;
        ptr = (void *)item + sizeof(my_mem_item_head_t);
        list->head.next = item->next;
        list->free_pt(ptr);
        --list->count;
    }
    my_spinlock_unlock(&list->lock);
}

#define MY_4KALIGNED_MEM_ITEM_OFFSET                            \
    ((MY_MEM_4K_ALIGNMENT - 1) + (int64_t)sizeof(my_aligned_mem_item_header_t))


#define MY_4KALIGNED_MEM_ITEM_PTR(item)                         \
    ((void *)(((size_t)item + (size_t)MY_4KALIGNED_MEM_ITEM_OFFSET) & ~((size_t)MY_MEM_4K_ALIGNMENT - 1)))

typedef struct {
    my_spinlock_t                   lock;
    uint32_t                        space_size;
    int64_t                         water_level;
    int64_t                         count;
    my_aligned_mem_item_header_t    head;
    my_aligned_malloc_pt            malloc_pt;
    my_aligned_free_pt              free_pt;            
} my_4kaligned_mem_freelist_t;

static inline void
my_4kaligned_mem_freelist_init(my_4kaligned_mem_freelist_t *list,
    uint32_t space_size, int64_t water_level,
    my_aligned_malloc_pt malloc_pt,
    my_aligned_free_pt free_pt)
{
    if (!malloc_pt || !free_pt) {
        assert(!"malloc or free handler is NULL");
    }

    my_spinlock_init(&list->lock);
    list->space_size = space_size;
    list->water_level = water_level;
    list->count = 0;
    list->head.magic = MY_ALIGNED_MEM_ITEM_MAGIC;
    list->head.size = space_size;
    list->head.next = NULL;
    list->head.ptr = NULL;
    list->malloc_pt = malloc_pt;
    list->free_pt = free_pt;
}

static inline void *
my_4kaligned_mem_freelist_alloc(my_4kaligned_mem_freelist_t *list)
{
    void *ptr                           = NULL;
    my_aligned_mem_item_header_t *item  = NULL;

    my_spinlock_lock(&list->lock);
    if (list->count > 0) {
        item = list->head.next;
        list->head.next = item->next;
        --list->count;
    }
    my_spinlock_unlock(&list->lock);

    if (item != NULL) {
        item->next = NULL;
        ptr = MY_4KALIGNED_MEM_ITEM_PTR(item);
        return ptr;
    }

    ptr = list->malloc_pt(MY_MEM_4K_ALIGNMENT, list->space_size);
    return ptr;
}

static inline void *
my_4kaligned_mem_freelist_calloc(my_4kaligned_mem_freelist_t *list)
{
    void *ptr = my_4kaligned_mem_freelist_alloc(list);
    bzero(ptr, (size_t)list->space_size);
    return ptr;
}

static inline void
my_4kaligned_mem_freelist_free(my_4kaligned_mem_freelist_t *list, void *ptr)
{
    void *p1 = ((void **)ptr)[-1];
    my_aligned_mem_item_header_t *item = (my_aligned_mem_item_header_t *)p1;

    if (item->magic != MY_ALIGNED_MEM_ITEM_MAGIC) {
        assert(!"unexcept magic number");
    }

    if (item->size != list->space_size) {
        assert(!"unexcept stats");
    }

    my_spinlock_lock(&list->lock);
    if (list->count < list->water_level) {
        item->next = list->head.next;
        list->head.next = item;
        item = NULL;
        ++list->count;
    }
    my_spinlock_unlock(&list->lock);

    if (item == NULL) {
        return;
    }

    list->free_pt(ptr);
}

static inline void
my_4kaligned_mem_freelist_clean(my_4kaligned_mem_freelist_t *list)
{
    my_aligned_mem_item_header_t    *item;
    void                            *ptr;

    my_spinlock_lock(&list->lock);
    while(list->head.next != NULL) {
        item = list->head.next;
        list->head.next = item->next;
        ptr = MY_4KALIGNED_MEM_ITEM_PTR(item);
        list->free_pt(ptr);
        --list->count;
    }
    my_spinlock_unlock(&list->lock);
}


typedef enum {
    MEM_FREELIST_MODE_32             = 0,
    MEM_FREELIST_MODE_64             = 1,
    MEM_FREELIST_MODE_128            = 2,
    MEM_FREELIST_MODE_256            = 3,
    MEM_FREELIST_MODE_512            = 4,
    MEM_FREELIST_MODE_1K             = 5,
    MEM_FREELIST_MODE_2K             = 6,
    MEM_FREELIST_MODE_4K             = 7,
    MEM_FREELIST_MODE_8K             = 8,
    MEM_FREELIST_MODE_16K            = 9,
    MEM_FREELIST_MODE_32K            = 10,
    MEM_FREELIST_MODE_64K            = 11,
    MEM_FREELIST_MODE_128K           = 12,
    MEM_FREELIST_MODE_256K           = 13,
    MEM_FREELIST_MODE_END            = 14
} my_mem_freelist_modle_t;

static inline uint32_t
my_mem_freelist_mode_to_size(int mode)
{
    return (uint32_t)(32 * (1 << mode));
}

typedef struct {
    my_mem_freelist_t       lists[MEM_FREELIST_MODE_END];
    unsigned                inited:1;
} my_mem_pool_t;

static inline void
my_mem_pool_init(my_mem_pool_t *pool, int64_t water_level,
    my_malloc_pt malloc_pt, my_free_pt free_pt)
{
    for (int i = 0; i < MEM_FREELIST_MODE_END; i++) {
        my_mem_freelist_init(&pool->lists[i],
            my_mem_freelist_mode_to_size(i),
            water_level, malloc_pt, free_pt);
    }

    pool->inited = 1;
}

static inline void *
my_mem_pool_alloc(my_mem_pool_t *pool, int64_t size)
{
    int mode = -1;
    for (int i = 0; i < MEM_FREELIST_MODE_END; i++) {
        if (size <= my_mem_freelist_mode_to_size(i)) {
            mode = i;
            break;
        }
    }

    if (mode == -1 || !pool->inited) {
        return my_mem_raw_alloc(size);
    }

    return my_mem_freelist_alloc(&pool->lists[mode]);
}

static inline void *
my_mem_pool_calloc(my_mem_pool_t *pool, int64_t size)
{
    void *ptr = my_mem_pool_alloc(pool, size);
    bzero(ptr, (size_t)size);
    return ptr;
}

static inline void
my_mem_pool_free(my_mem_pool_t *pool, void *ptr)
{
    my_mem_item_head_t *item;
    item = ptr - sizeof(my_mem_item_head_t);

    if (item->magic == MY_MEM_ITEM_RAW_MAGIC ||
        !pool->inited) {
        return my_mem_free(ptr);
    }

    if (item->magic != MY_MEM_ITEM_MAGIC) {
        assert(!"unexcept magic number");
    }

    int mode = -1;
    for (int i = 0; i < MEM_FREELIST_MODE_END; i++) {
        if (item->size <= my_mem_freelist_mode_to_size(i)) {
            mode = i;
            break;
        }
    }

    if (mode == -1 || !pool->inited) {
        assert(!"unexcept status");
    }

    return my_mem_freelist_free(&pool->lists[mode], ptr);
}

static inline void
my_mem_pool_clean(my_mem_pool_t *pool)
{
    for (int i = 0; i < MEM_FREELIST_MODE_END; i++) {
        my_mem_freelist_clean(&pool->lists[i]);
    }
}

// support 4KByte ~ 16MByte mem space aligned alloc
typedef enum {
    MEM_4KALIGEND_FREELIST_MODE_4K              = 0,
    MEM_4KALIGEND_FREELIST_MODE_8K              = 1,
    MEM_4KALIGEND_FREELIST_MODE_16K             = 2,
    MEM_4KALIGEND_FREELIST_MODE_32K             = 3,
    MEM_4KALIGEND_FREELIST_MODE_64K             = 4,
    MEM_4KALIGEND_FREELIST_MODE_128K            = 5,
    MEM_4KALIGEND_FREELIST_MODE_256K            = 6,
    MEM_4KALIGEND_FREELIST_MODE_512K            = 7,
    MEM_4KALIGEND_FREELIST_MODE_1M              = 9,
    MEM_4KALIGEND_FREELIST_MODE_2M              = 10,
    MEM_4KALIGEND_FREELIST_MODE_4M              = 11,
    MEM_4KALIGEND_FREELIST_MODE_8M              = 12,
    MEM_4KALIGEND_FREELIST_MODE_16M             = 13,
    MEM_4KALIGEND_FREELIST_MODE_END             = 14
} my_4kaligned_mem_freelist_modle_t;

static inline uint32_t
my_4kaligned_mem_freelist_mode_to_size(int mode)
{
    return (uint32_t)((4 * 1024) * (1 << mode));
}

typedef struct {
    my_4kaligned_mem_freelist_t                 lists[MEM_4KALIGEND_FREELIST_MODE_END];
    unsigned                                    inited:1;
} my_4kaligned_mem_pool_t;

static inline void
my_4kaligned_mem_pool_init(my_4kaligned_mem_pool_t *pool,
    int64_t water_level,
    my_aligned_malloc_pt malloc_pt,
    my_aligned_free_pt free_pt)
{
    for (int i = 0; i < MEM_4KALIGEND_FREELIST_MODE_END; i++) {
        my_4kaligned_mem_freelist_init(&pool->lists[i],
            my_4kaligned_mem_freelist_mode_to_size(i),
            water_level, malloc_pt, free_pt);
    }

    pool->inited = 1;
}

static inline void *
my_4kaligned_mem_pool_alloc(my_4kaligned_mem_pool_t *pool, int64_t size)
{
    int mode = -1;
    for (int i = 0; i < MEM_4KALIGEND_FREELIST_MODE_END; i++) {
        if (size <= my_4kaligned_mem_freelist_mode_to_size(i)) {
            mode = i;
            break;
        }
    }

    if (mode == -1 || !pool->inited) {
        return my_aligned_raw_malloc(MY_MEM_4K_ALIGNMENT, size);
    }

    return my_4kaligned_mem_freelist_alloc(&pool->lists[mode]);
}

static inline void
my_4kaligned_mem_pool_free(my_4kaligned_mem_pool_t *pool, void *ptr)
{
    void *p1 = ((void **)ptr)[-1];
    my_aligned_mem_item_header_t *item = p1;

    if (item->magic == MY_ALIGNED_MEM_ITEM_RAW_MAGIC ||
        !pool->inited) {
        return my_aligned_free(ptr);
    }

    if (item->magic != MY_ALIGNED_MEM_ITEM_MAGIC) {
        assert(!"unexcept magic number");
    }

    int mode = -1;
    for (int i = 0; i < MEM_4KALIGEND_FREELIST_MODE_END; i++) {
        if (item->size <= my_4kaligned_mem_freelist_mode_to_size(i)) {
            mode = i;
            break;
        }
    }

    if (mode == -1) {
        assert(!"unexcept status");
    }

    return my_4kaligned_mem_freelist_free(&pool->lists[mode], ptr);
}

static inline void
my_4kaligned_mem_pool_clean(my_4kaligned_mem_pool_t *pool)
{
    for (int i = 0; i < MEM_4KALIGEND_FREELIST_MODE_END; i++) {
        my_4kaligned_mem_freelist_clean(&pool->lists[i]);
    }
}

extern my_mem_pool_t                g_mem_pool;
extern my_4kaligned_mem_pool_t      g_4kaligned_mem_pool;